import os
import tkinter as tk
from tkinter import messagebox
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from PIL import Image
import xml.etree.ElementTree as ET
import numpy as np

image_type = "png"


class VOCViewer:
    def __init__(self, root, image_dir, ann_dir):
        self.root = root
        self.image_dir = image_dir
        self.ann_dir = ann_dir
        self.image_ids = [f.split(".xml")[0] for f in os.listdir(ann_dir) if f.endswith(".xml")]
        self.current_index = 0
        self.max_index = len(self.image_ids) - 1

        if not self.image_ids:
            messagebox.showerror("错误", "未找到任何标注文件！")
            self.root.quit()
            return

        # GUI 布局
        self.root.title("VOC 数据集查看器")

        # 绑定窗口关闭事件
        self.root.protocol("WM_DELETE_WINDOW", self.close)

        # 索引输入框
        self.label = tk.Label(root, text=f"图像索引 (0-{self.max_index}):")
        self.label.pack(pady=5)

        self.index_entry = tk.Entry(root)
        self.index_entry.pack(pady=5)
        self.index_entry.insert(0, "0")

        # 按钮框架
        self.button_frame = tk.Frame(root)
        self.button_frame.pack(pady=5)

        self.prev_button = tk.Button(self.button_frame, text="上一张", command=self.show_prev)
        self.prev_button.pack(side=tk.LEFT, padx=5)

        self.next_button = tk.Button(self.button_frame, text="下一张", command=self.show_next)
        self.next_button.pack(side=tk.LEFT, padx=5)

        self.show_button = tk.Button(self.button_frame, text="显示", command=self.show_image)
        self.show_button.pack(side=tk.LEFT, padx=5)

        # Matplotlib 画布
        self.fig, self.ax = plt.subplots()
        self.canvas = FigureCanvasTkAgg(self.fig, master=root)
        self.canvas.get_tk_widget().pack(pady=10)

        # 初始化显示第一张图像
        self.show_image()

    def load_voc_image_and_annotations(self, image_id):
        """加载图像和标注"""
        image_path = os.path.join(self.image_dir, f"{image_id}.{image_type}")
        ann_path = os.path.join(self.ann_dir, f"{image_id}.xml")

        # 加载图像
        img = Image.open(image_path)
        img = np.array(img)

        # 解析 XML
        tree = ET.parse(ann_path)
        root = tree.getroot()
        bboxes = []
        labels = []

        for obj in root.findall("object"):
            label = obj.find("name").text
            bbox = obj.find("bndbox")
            xmin = float(bbox.find("xmin").text)
            ymin = float(bbox.find("ymin").text)
            xmax = float(bbox.find("xmax").text)
            ymax = float(bbox.find("ymax").text)
            bboxes.append([xmin, ymin, xmax, ymax])
            labels.append(label)

        return img, bboxes, labels

    def plot_image_with_bboxes(self, img, bboxes, labels):
        """绘制图像和边界框"""
        self.ax.clear()
        self.ax.imshow(img)

        # 使用 tab10 颜色映射，规范化索引到 [0, 1]
        colors = plt.colormaps["tab10"]
        for i, (bbox, label) in enumerate(zip(bboxes, labels)):
            xmin, ymin, xmax, ymax = bbox
            rect = plt.Rectangle(
                (xmin, ymin), xmax - xmin, ymax - ymin, linewidth=2, edgecolor=colors(i % 10 / 9), facecolor="none"
            )
            self.ax.add_patch(rect)
            self.ax.text(xmin, ymin - 10, label, color=colors(i % 10 / 9), fontsize=10, weight="bold")

        self.ax.set_axis_off()
        self.ax.set_title(f"Image {self.current_index}/{self.max_index}")
        self.canvas.draw()

    def show_image(self):
        """显示指定索引的图像"""
        try:
            idx = int(self.index_entry.get())
            if 0 <= idx <= self.max_index:
                self.current_index = idx
                image_id = self.image_ids[self.current_index]
                img, bboxes, labels = self.load_voc_image_and_annotations(image_id)
                self.plot_image_with_bboxes(img, bboxes, labels)
            else:
                messagebox.showerror("错误", f"索引超出范围！请输入 0 到 {self.max_index}。")
        except ValueError:
            messagebox.showerror("错误", "请输入有效的数字索引！")
        except FileNotFoundError:
            messagebox.showerror("错误", f"未找到图像或标注文件：{image_id}")
        except Exception as e:
            messagebox.showerror("错误", f"发生错误：{e}")

    def show_prev(self):
        """显示上一张图像"""
        if self.current_index > 0:
            self.current_index -= 1
            self.index_entry.delete(0, tk.END)
            self.index_entry.insert(0, str(self.current_index))
            self.show_image()

    def show_next(self):
        """显示下一张图像"""
        if self.current_index < self.max_index:
            self.current_index += 1
            self.index_entry.delete(0, tk.END)
            self.index_entry.insert(0, str(self.current_index))
            self.show_image()

    def close(self):
        """处理窗口关闭事件，清理资源并退出"""
        plt.close(self.fig)  # 关闭 Matplotlib 图形
        self.root.destroy()  # 销毁 Tkinter 窗口
        self.root.quit()  # 退出 Tkinter 主循环


if __name__ == "__main__":
    # 设置 VOC 数据集路径
    ann_dir = "./result/2025dog/annotations"
    image_dir = "./result/2025dog/images"

    # 检查路径
    if not os.path.exists(image_dir) or not os.path.exists(ann_dir):
        print("请将路径设置为有效的 VOC 数据集路径！")
    else:
        root = tk.Tk()
        app = VOCViewer(root, image_dir, ann_dir)
        root.mainloop()
